import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation


def heart_function(t):
    """参数方程表示的爱心曲线"""
    x = 16 * np.sin(t) ** 3
    y = 13 * np.cos(t) - 5 * np.cos(2 * t) - 2 * np.cos(3 * t) - np.cos(4 * t)
    return x, y


def animate_heart(save_path=None):
    """创建爱心动画，可选择保存为文件"""
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.set_aspect('equal')
    ax.axis('off')
    line, = ax.plot([], [], color='red', linewidth=2)

    # 生成参数t的取值范围
    t = np.linspace(0, 2 * np.pi, 1000)

    def init():
        line.set_data([], [])
        return line,

    def update(frame):
        # 动态更新显示的点数，实现绘制动画效果
        x, y = heart_function(t[:frame])
        line.set_data(x, y)
        return line,

    # 创建动画
    ani = FuncAnimation(fig, update, frames=len(t), init_func=init,
                        interval=20, blit=True)

    plt.tight_layout()

    if save_path:
        ani.save(save_path, writer='pillow', fps=30)  # 保存为GIF
        print(f"动画已保存到: {save_path}")
    else:
        plt.show()

    return ani


# 调用函数创建动画或保存动画
if __name__ == "__main__":
    # 直接显示动画
    animate_heart()

    # 或者保存为GIF文件(取消下面一行的注释)
    # animate_heart('heart_animation.gif')